#include "flv_tag.hpp"
using namespace mms;

int32_t AudioTagHeader::decode(const uint8_t *data, size_t len)
{
    const uint8_t *data_start = data;
    if (len < 1)
    {
        return -1;
    }
    sound_type = (SoundType)((*data) & 0x01);
    sound_size = (SoundSize)(((*data) >> 1) & 0x01);
    sound_rate = (SoundRate)(((*data) >> 2) & 0x03);
    sound_format = (SoundFormat)((*data >> 4) & 0x0f);

    len--;
    data++;

    if (sound_format == AAC)
    {
        if (len < 1)
        {
            return -2;
        }
        aac_packet_type = (AACPacketType)(*data);
        len--;
        data++;
    }
    return data - data_start;
}

int32_t AudioTagHeader::encode(uint8_t *data, size_t len)
{
    auto data_start = data;
    if (len < 1)
    {
        return -1;
    }
    data[0] = (sound_format << 4) | (sound_rate << 2) | (sound_size << 1) | sound_type;
    len--;
    data++;

    if (sound_format == AAC)
    {
        if (len < 1)
        {
            return -2;
        }
        data[0] = aac_packet_type;
        data++;
        len--;
    }

    return data - data_start;
}

int32_t AUDIODATA::decode(const uint8_t *data, size_t len)
{
    auto data_start = data;
    auto consumed = header.decode(data, len);
    if (consumed < 0)
    {
        return -1;
    }
    data += consumed;
    len -= consumed;
    payload = std::string_view((char *)data, len);
    data += len;
    return data - data_start;
}

int32_t AUDIODATA::encode(uint8_t *data, size_t len)
{
    auto data_start = data;
    auto consumed = header.encode(data, len);
    if (consumed < 0)
    {
        return -1;
    }
    data += consumed;
    len -= consumed;
    if (len < payload.size())
    {
        return -2;
    }

    if (payload.data() == (const char *)data)
    {
        data += payload.size();
        return data - data_start;
    }
    memcpy(data, payload.data(), payload.size());
    payload = std::string_view((char *)data, payload.size());
    data += payload.size();
    return data - data_start;
}

bool VideoTagHeader::is_seq_header()
{
    if (codec_id == AVC || codec_id == HEVC)
    {
        return avc_packet_type == AVCSequenceHeader;
    }
    else
    {
        return false;
    }

    return false;
}

int32_t VideoTagHeader::decode(const uint8_t *data, size_t len)
{
    auto data_start = data;
    if (len < 1)
    {
        return 0;
    }

    frame_type = (FrameType)((*data >> 4) & 0x07);
    codec_id = CodecID(((*data)) & 0x0f);
    len--;
    data++;

    if (codec_id == AVC || codec_id == HEVC)
    {
        if (len < 4)
        {
            return 0;
        }
        avc_packet_type = (AVCPacketType)(*data);
        data++;
        len--;

        composition_time = 0;
        uint8_t *p = (uint8_t *)&composition_time;
        p[0] = data[2];
        p[1] = data[1];
        p[2] = data[0];
        data += 3;
        len -= 3;
    }
    else
    {
        return -1;
    }

    size = data - data_start;
    return data - data_start;
}

int32_t VideoTagHeader::encode(uint8_t *data, size_t len)
{
    auto data_start = data;
    if (len < 1)
    {
        return 0;
    }

    data[0] = (frame_type << 4) | codec_id;
    len--;
    data++;
    if (codec_id == AVC || codec_id == HEVC)
    {
        if (len < 4)
        {
            return -2;
        }
        *data = avc_packet_type;
        data++;
        len--;

        uint8_t *p = (uint8_t *)&composition_time;
        data[0] = p[2];
        data[1] = p[1];
        data[2] = p[0];
        data += 3;
        len -= 3;
    }
    else
    {
        return -1;
    }

    return data - data_start;
}

int32_t VIDEODATA::decode(const uint8_t *data, size_t len)
{
    auto data_start = data;
    auto consumed = header.decode(data, len);
    if (consumed <= 0)
    {
        return consumed;
    }
    data += consumed;
    len -= consumed;
    payload = std::string_view((char *)data, len);
    data += len;
    return data - data_start;
}

int32_t VIDEODATA::encode(uint8_t *data, size_t len)
{
    auto data_start = data;
    auto consumed = header.encode(data, len);
    if (consumed < 0)
    {
        return -1;
    }
    data += consumed;
    len -= consumed;
    if (len < payload.size())
    {
        // spdlog::error("len:{}, payload:{}", len, payload.size());
        return -2;
    }

    memcpy(data, payload.data(), payload.size());
    payload = std::string_view((char *)data, payload.size());
    data += payload.size();
    return data - data_start;
}

int32_t FlvTag::decode(const uint8_t *data, size_t len)
{
    auto data_start = data;
    auto consumed = tag_header.decode(data, len);
    if (consumed < 0)
    {
        return -1;
    }
    else if (consumed == 0)
    {
        return 0;
    }

    data += consumed;
    len -= consumed;

    if (len < tag_header.data_size)
    {
        return 0;
    }

    if ((tag_header.tag_type & 0x1f) == FlvTagHeader::AudioTag)
    {
        tag_data = std::make_unique<AUDIODATA>();
        consumed = tag_data->decode(data, tag_header.data_size);
        if (consumed < 0)
        {
            return -2;
        }
    }
    else if ((tag_header.tag_type & 0x1f) == FlvTagHeader::VideoTag)
    {
        tag_data = std::make_unique<VIDEODATA>();
        consumed = tag_data->decode(data, tag_header.data_size);
        if (consumed < 0)
        {
            return -2;
        }

        if (consumed > 0)
        {
            VIDEODATA *video_data = (VIDEODATA *)tag_data.get();
            if (video_data->header.codec_id == VideoTagHeader::AVC)
            {
                video_data->set_pts(tag_header.timestamp + video_data->header.composition_time);
                video_data->set_dts(tag_header.timestamp);
            }
        }
    }
    else if ((tag_header.tag_type & 0x1f) == FlvTagHeader::ScriptTag)
    {
        tag_data = std::make_unique<SCRIPTDATA>();
        consumed = tag_data->decode(data, tag_header.data_size);
        if (consumed < 0)
        {
            return -2;
        }
    }
    else
    {
        return -3;
    }

    if (consumed == 0)
    {
        return 0;
    }
    data += consumed;
    return data - data_start;
}

int32_t FlvTag::decode_and_store(const uint8_t *data, size_t len)
{
    auto data_start = data;
    auto tag_header_consumed = tag_header.decode(data, len);
    if (tag_header_consumed < 0)
    {
        return -1;
    }
    else if (tag_header_consumed == 0)
    {
        return 0;
    }

    data += tag_header_consumed;
    len -= tag_header_consumed;

    if (len < tag_header.data_size)
    {
        return 0;
    }

    auto total_size = tag_header_consumed + tag_header.data_size;
    // 将数据存储起来
    raw_data_size = total_size;
    raw_data_buf = std::make_unique<uint8_t[]>(total_size);
    memcpy(raw_data_buf.get(), data_start, raw_data_size);
    // 解析TAGDATA字段
    if ((tag_header.tag_type & 0x1f) == FlvTagHeader::AudioTag)
    {
        tag_data = std::make_unique<AUDIODATA>();
        auto consumed = tag_data->decode(raw_data_buf.get() + tag_header_consumed, tag_header.data_size);
        if (consumed < 0)
        {
            return -2;
        }
    }
    else if ((tag_header.tag_type & 0x1f) == FlvTagHeader::VideoTag)
    {
        tag_data = std::make_unique<VIDEODATA>();
        auto consumed = tag_data->decode(raw_data_buf.get() + tag_header_consumed, tag_header.data_size);
        if (consumed < 0)
        {
            return -2;
        }

        if (consumed > 0)
        {
            VIDEODATA *video_data = (VIDEODATA *)tag_data.get();
            if (video_data->header.codec_id == VideoTagHeader::AVC)
            {
                video_data->set_pts(tag_header.timestamp + video_data->header.composition_time);
                video_data->set_dts(tag_header.timestamp);
            }
        }
    }
    else if ((tag_header.tag_type & 0x1f) == FlvTagHeader::ScriptTag)
    {
        tag_data = std::make_unique<SCRIPTDATA>();
        auto consumed = tag_data->decode(raw_data_buf.get() + tag_header_consumed, tag_header.data_size);
        if (consumed < 0)
        {
            return -2;
        }
    }
    else
    {
        return -3;
    }

    return total_size;
}

int32_t FlvTag::encode(uint8_t *data, size_t len)
{
    auto data_start = data;
    int32_t consumed = tag_header.encode(data, len);
    if (consumed < 0)
    {
        return -1;
    }
    data += consumed;
    len -= consumed;
    consumed = tag_data->encode(data, len);
    if (consumed < 0)
    {
        return -2;
    }
    data += consumed;
    return data - data_start;
}